package com.zhang.second.day03;

import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * @title: 自定义物理分区
 * @author: zhang
 * @date: 2022/1/30 11:34
 */
public class Example1 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);


        //keyBy之后相同key的数据一定发送到同一个物理分区，不同key的数据也可能会在同一个物理分区
        //keyBy只是为每个数据指定key，具体发送到哪个物理分区不确定
        env
                .fromElements(1, 2, 3, 4, 5, 6, 7, 8, 9, 10).setParallelism(1)
                .keyBy(r -> r % 2)
                .print("keyBy:").setParallelism(2);

        //自定义物理分区，为每个元素指定key的同时，明确指定对应key的数据发送到哪个物理分区
        env
                .fromElements(1, 2, 3, 4)
                .partitionCustom(
                        //指定将某个key的数据发送到哪一个物理分区
                        new Partitioner<Integer>() {
                            @Override
                            public int partition(Integer key, int numPartitions) {
                                //return key % numPartitions;
                                return key;
                            }
                        },
                        //为每个元素指定key
                        new KeySelector<Integer, Integer>() {
                            @Override
                            public Integer getKey(Integer value) throws Exception {
                                return value % 2;
                            }
                        }
                )
                .print().setParallelism(2);

        env.execute();
    }
}
